"""
author: liuxu
date: 2025/5/1
description: 此模块用于测试--omp-trace功能
"""
import os

import allure
import pytest
import re
from common.logger_ctrl import mylogger
from common.ssh_command import exec_cmd
from base_public.testbase import BasePublic
from setting import CPP_DIR, EXE_DIR

@allure.feature('omp tracing指令')
@pytest.mark.medium
class TestOmpTrace(BasePublic):
    """测试omp trace功能"""
    @pytest.mark.smoke
    @pytest.mark.high
    def test_omp_trace(self, test_env):
        """
        测试hipprof --omp-trace
        :return:
        """
        ompcase = self.compile_omp_case()
        ret, output = exec_cmd(f"hipprof --omp-trace {ompcase}", cwd=test_env)
        assert ret is True
        output_file_d = self.get_all_prof_files_info(output)
        omp_api_trace_info = self.get_trace_info_from_csv(output_file_d["omp_api_csv"]['abspath'])
        omp_kernel_trace_info = self.get_trace_info_from_csv(output_file_d["omp_kernel_csv"]['abspath'])

        # 检查调用的api和kernel次数是否符合预期
        expect_kernel = ["__nv_test_vec_mult__TARGET_F1L13_1_",
                         "__nv_test_vec_mult__TARGET_F1L25_2_",
                         "__nv_test_vec_mult__TARGET_F1L38_3_",
                         "__nv_test_vec_mult__TARGET_F1L51_4_"]
        assert omp_api_trace_info.get("Total", {}).get("Calls", 0) >= 10
        assert set(omp_kernel_trace_info.keys()) >= set(expect_kernel)

    @pytest.mark.high
    def test_omp_libc_hsa_trace(self, test_env):
        """
        测试hipprof --omp-trace --libc-trace --hsa-trace组合参数测试
        :return:
        """
        ompcase = self.compile_omp_case()
        ret, output = exec_cmd(f"hipprof --omp-trace --libc-trace --hsa-trace {ompcase}", cwd=test_env)
        assert ret is True
        output_file_d = self.get_all_prof_files_info(output)
        # 校验每个trace都有数据
        for csv_name in ["omp_api_csv", "omp_kernel_csv", "libc_api_csv", "hsa_api_csv"]:
            assert csv_name in output_file_d and len(self.get_trace_info_from_csv(output_file_d[csv_name]['abspath'])) >= 1

    # --------------------------------------------------------------------------------------------------------------
    # dtk-25.04.2开始测试omp&hip混合使用场景
    @pytest.mark.high
    def test_omp_hip_mixed_trace(self, test_env):
        """
        测试hipprof --omp-trace --hip-trace组合参数测试
        :return:
        """
        # 编译用例
        makefile_path = os.path.join(CPP_DIR, "25.04.2/omp_hip_mixed/Makefile")
        build_path = os.path.join(test_env, "build")
        assert self.create_directory(build_path) is True  # 创建build构建用例目录
        ret0, output0 = exec_cmd(f"make -f {makefile_path}", cwd=build_path)  # 工作目录得改到build里

        # 检测是否生成可执行文件
        assert ret0 is True, f"Make failed"
        exe_path = os.path.join(build_path, "test_hip_omp_mixed_program")
        assert self.exist_file(exe_path) is True, f"No executable file ({exe_path}) was found"

        # 执行测试
        ret, output = exec_cmd(f"hipprof --omp-trace --hip-trace {exe_path}", cwd=test_env)
        assert ret is True
        output_file_d = self.get_all_prof_files_info(output)
        # 校验每个trace都有数据
        for csv_name in ["omp_api_csv", "omp_kernel_csv", "hip_api_csv", "hip_kernel_csv"]:
            assert csv_name in output_file_d and len(self.get_trace_info_from_csv(output_file_d[csv_name]['abspath'])) >= 1

    @pytest.mark.high
    def test_omp_hip_libc_hsa_mixed_trace(self, test_env):
        """
        测试hipprof --omp-trace --hip-trace --libc-trace --hsa-trace组合参数测试
        :return:
        """
        # 编译用例
        makefile_path = os.path.join(CPP_DIR, "25.04.2/omp_hip_mixed/Makefile")
        build_path = os.path.join(test_env, "build")
        assert self.create_directory(build_path) is True  # 创建build构建用例目录
        ret0, output0 = exec_cmd(f"make -f {makefile_path}", cwd=build_path)  # 工作目录得改到build里

        # 检测是否生成可执行文件
        assert ret0 is True, f"Make failed"
        exe_path = os.path.join(build_path, "test_hip_omp_mixed_program")
        assert self.exist_file(exe_path) is True, f"No executable file ({exe_path}) was found"

        # 执行测试
        ret, output = exec_cmd(f"hipprof --omp-trace --hip-trace --libc-trace --hsa-trace --trace-args {exe_path}", cwd=test_env)
        assert ret is True
        output_file_d = self.get_all_prof_files_info(output)
        # 校验每个trace都有数据
        for csv_name in ["omp_api_csv", "omp_kernel_csv", "hip_api_csv", "hip_kernel_csv", "libc_api_csv", "hsa_api_csv"]:
            assert csv_name in output_file_d and len(self.get_trace_info_from_csv(output_file_d[csv_name]['abspath'])) >= 1

        #TODO 校验每个trace类型的函数都有参args
